import os
import numpy as np
import pandas as pd
from filterpy.kalman import KalmanFilter

def create_kalman_filter():
    kf = KalmanFilter(dim_x=4, dim_z=2)
    # 定义状态转移矩阵（描述系统如何从一个状态转移到下一个状态）
    kf.F = np.array([[1, 0, 1, 0],
                     [0, 1, 0, 1],
                     [0, 0, 1, 0],
                     [0, 0, 0, 1]])
    # 定义观测矩阵（描述如何从状态获得观测）
    kf.H = np.array([[1, 0, 0, 0],
                     [0, 1, 0, 0]])
    # 定义过程噪声协方差矩阵（描述系统的不确定性）
    kf.Q = np.array([[1, 0, 0, 0],
                     [0, 1, 0, 0],
                     [0, 0, 1, 0],
                     [0, 0, 0, 1]])
    # 定义观测噪声协方差矩阵（描述观测的不确定性）
    kf.R = np.array([[100, 0],
                     [0, 100]])
    # 定义初始状态估计
    kf.x = np.array([0, 0, 0, 0])
    # 定义初始协方差估计
    kf.P = np.eye(4)
    return kf

def execute_kalman_filter(path: np.ndarray) -> np.ndarray:
    result = path - path[0]
    for i in range(4):
        kf = create_kalman_filter()
        path = result
        result = []
        for measurement in path:
            kf.predict()
            kf.update(measurement)
            result.append(kf.x[:2])
    return np.array(result)

if __name__ == "__main__":
    df = pd.read_csv(os.path.join("data","100063.csv"))
    df = df[df.OBJECT_TYPE == 'AGENT']
    path = df[['X', 'Y']].values
    ft = execute_kalman_filter(path)
    for p in ft:
        print(f"{p[0]:.7f}, {p[1]:.7f}")